import logging
import socket
import sys

import dpkt
from scapy.layers.inet import *
from scapy.layers.inet6 import *

from app.utils import nethash, searcher, utc_stamp_to_utc

LOGGER = logging.getLogger(__name__)

class PacketRecord(object):
    """
    An object that represents a packet
    """

    def __init__(self, packet, mode = "b", packet_type="scapy"):
        """
        Parameters:
            timestamp: when packet was recorded
            packet: dpkt object
            mode: b (bidirectional) or u (unidirectional). Used for
            hash calculation
        """
        # *** Initialise packet variables:
        self.flow_hash = 0
        self.timestamp = packet.time
        # self.length = len(packet)
        self.ip_src = 0
        self.ip_dst = 0
        self.proto = 0
        self.tp_src = 0
        self.tp_dst = 0
        self.tp_flags = 0
        self.tp_seq_src = 0
        self.tp_seq_dst = 0
        self.region_src = ""
        self.region_dst = ""
        self.ingested = False
        self.is_ipv4 = True

        # packet handle by dpkt
        if packet_type == "dpkt":
            try:
                # Read packet into dpkt to parse headers:
                eth = dpkt.ethernet.Ethernet(packet)
            except:
                # Skip Packet if unable to parse:
                LOGGER.error("failed to unpack packet, skipping...")
                return

            # Get the IP packet
            ip = eth.data

            # Get the length of IPv4 packet:
            if isinstance(eth.data, dpkt.ip.IP):
                self.length = ip.len
            # Get the length of IPv6 packet:
            elif isinstance(eth.data, dpkt.ip6.IP6):
                self.length = len(ip.data)
                self.is_ipv4 = False
            # Ignore if non-IP packet:
            else:
                return

            # Handle IPv4 and IPv6:
            try:
                self.ip_src = socket.inet_ntop(socket.AF_INET, ip.src)
                self.ip_dst = socket.inet_ntop(socket.AF_INET, ip.dst)
            except ValueError:
                self.ip_src = socket.inet_ntop(socket.AF_INET6, ip.src)
                self.ip_dst = socket.inet_ntop(socket.AF_INET6, ip.dst)
            # Transport layer:
            self.proto = ip.p
            if ip.p == 6:
                # TCP
                tcp = ip.data
                self.tp_src = tcp.sport
                self.tp_dst = tcp.dport
                self.tp_flags = tcp.flags
                self.tp_seq_src = tcp.seq
                self.tp_seq_dst = tcp.ack
            elif ip.p == 17:
                # UDP
                udp = ip.data
                self.tp_src = udp.sport
                self.tp_dst = udp.dport
                self.tp_flags = ""
                self.tp_seq_src = 0
                self.tp_seq_dst = 0
            else:
                # Not a transport layer that we understand, keep going:
                pass

        # packet handle by scapy
        if packet_type == "scapy":
            if packet.haslayer(IP):
                self.length = packet[IP].len
                self.ip_src = packet[IP].src
                self.ip_dst = packet[IP].dst
                self.proto = packet[IP].proto
            elif packet.haslayer(IPv6):
                self.is_ipv4 = False
                self.length = packet[IPv6].plen
                self.ip_src = packet[IPv6].src
                self.ip_dst = packet[IPv6].dst
            else:
                return
            # Transport layer:
            if self.proto == 6:
                # TCP
                self.tp_src = packet[TCP].sport
                self.tp_dst = packet[TCP].dport
                self.tp_flags = packet[TCP].flags
                self.tp_seq_src = packet[TCP].seq
                self.tp_seq_dst = packet[TCP].ack
            elif self.proto == 17:
                # UDP
                self.tp_src = packet[UDP].sport
                self.tp_dst = packet[UDP].dport
                self.tp_flags = ""
                self.tp_seq_src = 0
                self.tp_seq_dst = 0
            else:
                # Not a transport layer that we understand, keep going:
                pass

        if mode == 'b':
            if self.proto == 6 or self.proto == 17:
                # Generate a directional 5-tuple flow_hash:
                self.flow_hash = nethash.hash_b5((self.ip_src,
                                                  self.ip_dst, self.proto, self.tp_src,
                                                  self.tp_dst))
            else:
                # Generate a directional 3-tuple flow_hash:
                self.flow_hash = nethash.hash_b3((self.ip_src,
                                                  self.ip_dst, self.proto))
        elif mode == 'u':
            if self.proto == 6 or self.proto == 17:
                # Generate a directional 5-tuple flow_hash:
                self.flow_hash = nethash.hash_u5((self.ip_src,
                                                  self.ip_dst, self.proto, self.tp_src,
                                                  self.tp_dst))
            else:
                # Generate a directional 3-tuple flow_hash:
                self.flow_hash = nethash.hash_u3((self.ip_src,
                                                  self.ip_dst, self.proto))
        else:
            LOGGER.critical("unsupported mode=%s", mode)
            sys.exit()
        
        self.region_src = searcher.search(self.ip_src) if self.is_ipv4 else ""
        self.region_dst = searcher.search(self.ip_dst) if self.is_ipv4 else ""
        # Yay, packet has been ingested:
        self.ingested = True

    def tcp_fin(self):
        """
        Does the current packet have the TCP FIN flag set?
        """
        return self.tp_flags & dpkt.tcp.TH_FIN != 0

    def tcp_syn(self):
        """
        Does the current packet have the TCP SYN flag set?
        """
        return self.tp_flags & dpkt.tcp.TH_SYN != 0

    def tcp_rst(self):
        """
        Does the current packet have the TCP RST flag set?
        """
        return self.tp_flags & dpkt.tcp.TH_RST != 0

    def tcp_psh(self):
        """
        Does the current packet have the TCP PSH flag set?
        """
        return self.tp_flags & dpkt.tcp.TH_PUSH != 0

    def tcp_ack(self):
        """
        Does the current packet have the TCP ACK flag set?
        """
        return self.tp_flags & dpkt.tcp.TH_ACK != 0

    def tcp_urg(self):
        """
        Does the current packet have the TCP URG flag set?
        """
        return self.tp_flags & dpkt.tcp.TH_URG != 0

    def tcp_ece(self):
        """
        Does the current packet have the TCP ECE flag set?
        """
        return self.tp_flags & dpkt.tcp.TH_ECE != 0

    def tcp_cwr(self):
        """
        Does the current packet have the TCP CWR flag set?
        """
        return self.tp_flags & dpkt.tcp.TH_CWR != 0

    def todict(self):
        return {
            'Time': utc_stamp_to_utc(self.timestamp),
            'Source IP': self.ip_src,
            'Source Port': self.tp_src,
            'Source Region': self.region_src,
            'Destination IP': self.ip_dst,
            'Destination Port': self.tp_dst,
            'Destination Region': self.region_dst
        }

